clear
load('counter.mat')

T_ftr = 2;
shock_size_ftr = 0.2;

intra_id = 0;

%%% Baseline: Everything stays the same

pop_ftr0 = zeros(N,T_ftr+1);
pop_growth_ftr0 = zeros(N,1);

pop_ftr0(:,1) = pop_dyn(:,end);

alpha_ftr = repmat(alpha_dyn(:,end),1,T_ftr+1);

lambda_ftr0 = squeeze(lambda_dyn(:,:,end));
L_k_ftr = L_k_dyn(:,end); L_y_ftr = squeeze(L_y_dyn(:,:,end));
migration_exposure_ftr = migration_exposure_dyn(:,:,end);

for t = 1+1:1+T_ftr
    
    L_k_ftr_aux = L_k_ftr;
    
    [pop_ftr0(:,t),lambda_ftr0,L_k_ftr,L_y_ftr,L_o_ftr,pi_ftr] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_ftr,intra_id,lambda_ftr0,alpha_ftr(:,t),squeeze(epsilon_dyn(:,:,end)),...
        pop_ftr0(:,t-1),L_k_ftr,L_y_ftr,ubar_dyn(:,end),mu,theta,zeta,omega,fertility_dyn(end));

    prob_orig_given_dest_y_ftr = zeros(N,N);
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_ftr(i,j) = sum(pi_ftr(i,j,:))*L_k_ftr_aux(i) / sum(L_y_ftr(j,:));
        end
    end
    migration_exposure_ftr = prob_orig_given_dest_y_ftr.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_ftr.*(1-eye(N)),1);
    
end

pop_growth_ftr0(:,1) = log(pop_ftr0(:,end)) - log(pop_ftr0(:,1));

%%% First scenario: B2 (s=5) increases by 2*std

pop_ftr1 = zeros(N,T_ftr+1);
pop_growth_ftr1 = zeros(N,1);

pop_ftr1(:,1) = pop_dyn(:,end);

alpha_ftr = repmat(alpha_dyn(:,end),1,T_ftr+1);
alpha_ftr(5,2:end) = (1+shock_size_ftr)*alpha_dyn(5,end);

lambda_ftr1 = squeeze(lambda_dyn(:,:,end));
L_k_ftr = L_k_dyn(:,end); L_y_ftr = squeeze(L_y_dyn(:,:,end));
migration_exposure_ftr = migration_exposure_dyn(:,:,end);

for t = 1+1:1+T_ftr
    
    L_k_ftr_aux = L_k_ftr;
    
    [pop_ftr1(:,t),lambda_ftr1,L_k_ftr,L_y_ftr,L_o_ftr,pi_ftr] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_ftr,intra_id,lambda_ftr1,alpha_ftr(:,t),squeeze(epsilon_dyn(:,:,end)),...
        pop_ftr1(:,t-1),L_k_ftr,L_y_ftr,ubar_dyn(:,end),mu,theta,zeta,omega,fertility_dyn(end));
    
    prob_orig_given_dest_y_ftr = zeros(N,N);
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_ftr(i,j) = sum(pi_ftr(i,j,:))*L_k_ftr_aux(i) / sum(L_y_ftr(j,:));
        end
    end
    migration_exposure_ftr = prob_orig_given_dest_y_ftr.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_ftr.*(1-eye(N)),1);

end

pop_growth_ftr1(:,1) = log(pop_ftr1(:,end)) - log(pop_ftr1(:,1));

%%% Second scenario: Composite diffusion frictions (diffusion_frictions)
%%% decline by 50% between s=5 and s=10, s=11

pop_ftr2 = zeros(N,T_ftr+1);
pop_growth_ftr2 = zeros(N,1);

pop_ftr2(:,1) = pop_dyn(:,end);

alpha_ftr = repmat(alpha_dyn(:,end),1,T_ftr+1);

% d_know^theta to and from 5 and (10,11) goes down by 50%:
% d_know_ftr2^theta = 0.5*d_know^theta => d_know_ftr2 = (0.5^(1/theta))*d_know
d_know_ftr2 = d_know;
d_know_ftr2(5,10) = (0.5^(1/theta))*d_know(5,10);
d_know_ftr2(5,11) = (0.5^(1/theta))*d_know(5,11);
d_know_ftr2(10,5) = (0.5^(1/theta))*d_know(10,5);
d_know_ftr2(11,5) = (0.5^(1/theta))*d_know(11,5);

lambda_ftr2 = squeeze(lambda_dyn(:,:,end));
L_k_ftr = L_k_dyn(:,end); L_y_ftr = squeeze(L_y_dyn(:,:,end));
migration_exposure_ftr = migration_exposure_dyn(:,:,end);

for t = 1+1:1+T_ftr
    
    L_k_ftr_aux = L_k_ftr;
    
    [pop_ftr2(:,t),lambda_ftr2,L_k_ftr,L_y_ftr,L_o_ftr,pi_ftr] = counter_execute(N,S,d_geo,d_know_ftr2,d_geo_distance,d_migr_const_aux,migration_exposure_ftr,intra_id,lambda_ftr2,alpha_ftr(:,t),squeeze(epsilon_dyn(:,:,end)),...
        pop_ftr2(:,t-1),L_k_ftr,L_y_ftr,ubar_dyn(:,end),mu,theta,zeta,omega,fertility_dyn(end));
    
    prob_orig_given_dest_y_ftr = zeros(N,N);
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_ftr(i,j) = sum(pi_ftr(i,j,:))*L_k_ftr_aux(i) / sum(L_y_ftr(j,:));
        end
    end
    migration_exposure_ftr = prob_orig_given_dest_y_ftr.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_ftr.*(1-eye(N)),1);

end

pop_growth_ftr2(:,1) = log(pop_ftr2(:,end)) - log(pop_ftr2(:,1));

%%% Third scenario: A3 (s=3) increases by 2*std

pop_ftr3 = zeros(N,T_ftr+1);
pop_growth_ftr3 = zeros(N,1);

pop_ftr3(:,1) = pop_dyn(:,end);

alpha_ftr = repmat(alpha_dyn(:,end),1,T_ftr+1);
alpha_ftr(3,2:end) = (1+shock_size_ftr)*alpha_dyn(3,end);

lambda_ftr3 = squeeze(lambda_dyn(:,:,end));
L_k_ftr = L_k_dyn(:,end); L_y_ftr = squeeze(L_y_dyn(:,:,end));
migration_exposure_ftr = migration_exposure_dyn(:,:,end);

for t = 1+1:1+T_ftr
    
    L_k_ftr_aux = L_k_ftr;

    [pop_ftr3(:,t),lambda_ftr3,L_k_ftr,L_y_ftr,L_o_ftr,pi_ftr] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_ftr,intra_id,lambda_ftr3,alpha_ftr(:,t),squeeze(epsilon_dyn(:,:,end)),...
        pop_ftr3(:,t-1),L_k_ftr,L_y_ftr,ubar_dyn(:,end),mu,theta,zeta,omega,fertility_dyn(end));
    
    prob_orig_given_dest_y_ftr = zeros(N,N);
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_ftr(i,j) = sum(pi_ftr(i,j,:))*L_k_ftr_aux(i) / sum(L_y_ftr(j,:));
        end
    end
    migration_exposure_ftr = prob_orig_given_dest_y_ftr.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_ftr.*(1-eye(N)),1);

end

pop_growth_ftr3(:,1) = log(pop_ftr3(:,end)) - log(pop_ftr3(:,1));

%%% Fourth scenario: A1 (s=1) increases by 2*std

pop_ftr4 = zeros(N,T_ftr+1);
pop_growth_ftr4 = zeros(N,1);

pop_ftr4(:,1) = pop_dyn(:,end);

alpha_ftr = repmat(alpha_dyn(:,end),1,T_ftr+1);
alpha_ftr(1,2:end) = (1+shock_size_ftr)*alpha_dyn(1,end);

lambda_ftr4 = squeeze(lambda_dyn(:,:,end));
L_k_ftr = L_k_dyn(:,end); L_y_ftr = squeeze(L_y_dyn(:,:,end));
migration_exposure_ftr = migration_exposure_dyn(:,:,end);

for t = 1+1:1+T_ftr
    
    L_k_ftr_aux = L_k_ftr;

    [pop_ftr4(:,t),lambda_ftr4,L_k_ftr,L_y_ftr,L_o_ftr,pi_ftr] = counter_execute(N,S,d_geo,d_know,d_geo_distance,d_migr_const_aux,migration_exposure_ftr,intra_id,lambda_ftr4,alpha_ftr(:,t),squeeze(epsilon_dyn(:,:,end)),...
        pop_ftr4(:,t-1),L_k_ftr,L_y_ftr,ubar_dyn(:,end),mu,theta,zeta,omega,fertility_dyn(end));
    
    prob_orig_given_dest_y_ftr = zeros(N,N);
    for i=1:N
        for j=1:N
            prob_orig_given_dest_y_ftr(i,j) = sum(pi_ftr(i,j,:))*L_k_ftr_aux(i) / sum(L_y_ftr(j,:));
        end
    end
    migration_exposure_ftr = prob_orig_given_dest_y_ftr.*(1-eye(N)) ./ sum(prob_orig_given_dest_y_ftr.*(1-eye(N)),1);

end

pop_growth_ftr4(:,1) = log(pop_ftr4(:,end)) - log(pop_ftr4(:,1));

[ones(N,1) log(pop_dyn(:,end))] \ (pop_growth_ftr4-pop_growth_ftr0)

for_stata = [cz_list pop_growth_ftr1-pop_growth_ftr0, pop_growth_ftr2-pop_growth_ftr0, pop_growth_ftr3-pop_growth_ftr0, pop_growth_ftr4-pop_growth_ftr0];
csvwrite('results_model\future_scenarios_results.csv',for_stata)